import io
import logging
import time
from typing import Any, Iterable
from unittest.mock import patch

import pyzipper
import requests
from django.utils import timezone

from api_app.ingestors_manager.classes import Ingestor
from api_app.ingestors_manager.exceptions import IngestorRunException
from api_app.mixins import AbuseCHMixin
from tests.mock_utils import MockUpResponse, if_mock_connections

logger = logging.getLogger(__name__)


class MalwareBazaar(AbuseCHMixin, Ingestor):
    # API endpoint
    url: str
    # Download samples that are up to X hours old
    hours: int
    # Download samples from chosen signatures (aka malware families)
    signatures: str
    # Max number of results you want to display
    limit: int

    @classmethod
    def update(cls) -> bool:
        pass

    # retrieve information about the given signature
    def get_signature_information(self, signature):
        result = requests.post(
            self.url,
            data={"query": "get_siginfo", "signature": signature, "limit": self.limit},
            headers=self.authentication_header,
            timeout=30,
        )
        result.raise_for_status()
        content = result.json()
        logger.info(f"Malware bazaar data for signature {signature} is {content}")
        if content["query_status"] != "ok":
            raise IngestorRunException(
                f"Query status is invalid: {content['query_status']}"
            )
        if not isinstance(content["data"], list):
            raise IngestorRunException(f"Content {content} not expected")
        return content["data"]

    # extract file hashes per signature
    def get_recent_samples(self):
        hashes = set()
        current_time = timezone.now()
        for signature in self.signatures:
            data = self.get_signature_information(signature)
            hours_str = "hour" if self.hours == 1 else "hours"
            if len(data) > self.limit:
                logger.info(
                    f"{signature}: in the last {hours_str} there are "
                    f"more results than the limit {len(data)}/{self.limit}"
                )
            for elem in data:
                first_seen = timezone.make_aware(
                    timezone.datetime.strptime(elem["first_seen"], "%Y-%m-%d %H:%M:%S")
                )
                diff = int((current_time - first_seen).total_seconds()) // 3600
                if elem["signature"] == signature and diff <= self.hours:
                    hashes.add(elem["sha256_hash"])
        return hashes

    def download_sample(self, h):
        logger.info(f"Downloading sample {h}")
        sample_archive = requests.post(
            self.url,
            data={
                "query": "get_file",
                "sha256_hash": h,
            },
            headers=self.authentication_header,
            timeout=60,
        )
        sample_archive.raise_for_status()
        logger.info(f"Correctly downloaded sample {h}")
        logger.info("Sleeping for 1 second")
        time.sleep(1)
        with pyzipper.AESZipFile(io.BytesIO(sample_archive.content)) as zf:
            zf.setpassword(b"infected")
            files = zf.namelist()
            # expected only one file
            if files and len(files) == 1:
                return zf.read(files[0])

    def run(self) -> Iterable[Any]:
        hashes = self.get_recent_samples()
        hashes_len = len(hashes)
        # download sample and create new analysis
        for idx, h in enumerate(hashes):
            logger.info(f"Downloading sample {idx + 1}/{hashes_len}")
            sample = self.download_sample(h)
            yield sample

    @classmethod
    def _monkeypatch(cls):
        patches = [
            if_mock_connections(
                patch(
                    "requests.post",
                    return_value=MockUpResponse(
                        {
                            "query_status": "ok",
                            "data": [
                                {
                                    "sha256_hash": "c5c810beaf075f8fee52146b381b0f94a6"
                                    "e303fada3bce12bcc07fbfa07ba07e",
                                    "sha3_384_hash": "bdd25a594b5a5d8ab14b00c04ee75d6a"
                                    "476bf2a7df49223284eebfac82be107a"
                                    "b94ffaae294ef4cf0a1c23a206e1fbd9",
                                    "sha1_hash": "3fea40223c02a15678912a29147d2b32d05c"
                                    "46df",
                                    "md5_hash": "dc591fd6d108b50bd9aa1f3dce2f3fe4",
                                    "first_seen": "2024-04-11 12:35:10",
                                    "last_seen": None,
                                    "file_name": "17128389081d4616ae42b2693f5ea6783112"
                                    "f41cb2ee5184f49d983f8bf833df0b0e97b4"
                                    "29449.dat-decoded",
                                    "file_size": 240128,
                                    "file_type_mime": "application/x-dosexec",
                                    "file_type": "exe",
                                    "reporter": "abuse_ch",
                                    "anonymous": 0,
                                    "signature": "AgentTesla",
                                    "imphash": "f34d5f2d4577ed6d9ceec516c1f5a744",
                                    "tlsh": "T17534FD037E88EB15E5A87E3782EF6C2413B2B0C"
                                    "71633C60B6F49AF6518516426D7E72D",
                                    "telfhash": None,
                                    "gimphash": None,
                                    "ssdeep": "3072:z+ymieCL2QfOdb/TmqtbqRFP55EMX+CWQ:"
                                    "z+ymieCLPfOdbqq9qRFvXJW",
                                    "dhash_icon": None,
                                    "tags": ["AgentTesla", "base64-decoded", "exe"],
                                    "code_sign": [],
                                    "intelligence": {
                                        "clamav": None,
                                        "downloads": "338",
                                        "uploads": "1",
                                        "mail": None,
                                    },
                                }
                            ],
                        },
                        200,
                    ),
                ),
                patch(
                    "requests.post",
                    return_value=MockUpResponse(
                        {}, content=b"AgentTesla malware downloaded!", status_code=200
                    ),
                ),
            )
        ]
        return super()._monkeypatch(patches=patches)
